import os
import onnx
import numpy as np
from helper import create_onnx_session, load_onnx_model, check_and_save_onnx_model
from onnx import TensorProto, GraphProto
from collections import OrderedDict

def modify_onnx_model(onnx_model_name):
    ## load onnx model
    model = load_onnx_model(onnx_model_name)

    ## Delete the output graph node
    inputs = list()
    outputs = list()
    op_type = 'Pad'
    pad_name = 'Pad_1'
    model.graph.output.pop()
    for i in range(len(model.graph.node)):
        if model.graph.node[i].name == pad_name:
            model.graph.node[i].output.remove('y')
            model.graph.node[i].output.append(pad_name)

    ## Create the next 'Cast' node
    index = 2 ## The index is follow yourself onnx model
    op_type = 'Cast'
    op_name = '%s_%d' % (op_type, index)
    inputs = [pad_name]
    outputs = ['y']
    new_node = onnx.helper.make_node(
        op_type=op_type,
        name=op_name,
        inputs=inputs,
        outputs=outputs,
        to=getattr(TensorProto, 'FLOAT16'),
    )
    model.graph.node.append(new_node)

    ## Create output node
    output_value_info = onnx.helper.make_tensor_value_info("y", TensorProto.FLOAT16, shape=[1, 4])
    model.graph.output.append(output_value_info)
    ## Check that it works and re-save
    new_model_name = 'pad_cast.onnx'
    check_and_save_onnx_model(model, new_model_name)

    #########################################################
    ######## Check the modified onnx model file #############
    #########################################################
    ## Create onnx runtime
    sess = create_onnx_session(new_model_name)
    ## Generate input data
    x = np.random.randn(1, 2).astype(np.float32)
    pads = np.array([0, 1, 1, 2]).astype(np.int64)  # pad order [x1_begin, x2_begin, ..., x1_end, x2_end, ...]
    values = np.array([1.2]).astype(np.float32)
    print("input:", x)

    inputs = {"x":x, "pads":pads, "values":values}
    outputs = ['y']
    result = sess.run(outputs, inputs)
    print(result[0].shape)

def main():
    onnx_model_name = 'pad.onnx'
    modify_onnx_model(onnx_model_name)

if __name__ == "__main__":
    main()
